# -*- coding: utf-8 -*-
"""
Created on Thu Mar 21 14:26:55 2024

@author: Lenovo
"""


import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms

# 定义孪生神经网络模型
'''
class SiameseNetwork(nn.Module):
    def __init__(self):
        super(SiameseNetwork, self).__init__()
        # 定义神经网络结构，通常包括卷积层、池化层和全连接层

    def forward(self, input1, input2):
        output1 = self.siamese(input1)
        output2 = self.siamese(input2)
        return output1, output2
'''
'''
class SiameseNetwork(nn.Module):
    def __init__(self):
        super(SiameseNetwork, self).__init__()
        
        # 定义卷积层
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3),  # 输入通道为1（灰度图像），输出通道为32，卷积核大小为3x3
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 最大池化层，核大小为2x2
            nn.Conv2d(32, 64, kernel_size=3),  # 输入通道为32，输出通道为64，卷积核大小为3x3
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        
        # 定义全连接层
        self.fc = nn.Sequential(
            nn.Linear(64*22*22, 128),  # 计算公式：(图像大小/最大池化层步长)^2 * 最后一层卷积核数量 -> 22x22是根据输入图像大小和卷积、池化层设置计算得出的
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 8)  # 输出维度为8，代表特征向量的维度
        )

    def forward(self, input1, input2):
        output1 = self.cnn(input1)
        output1 = output1.view(output1.size()[0], -1)  # 将卷积层输出展平为一维向量
        output1 = self.fc(output1)

        output2 = self.cnn(input2)
        output2 = output2.view(output2.size()[0], -1)
        output2 = self.fc(output2)

        return output1, output2
'''
class SiameseNetwork(nn.Module):
    def __init__(self):
        super(SiameseNetwork, self).__init__()
        
        # 定义卷积层
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(32, 64, kernel_size=3),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        
        # 定义两个独立的全连接层
        self.fc1 = nn.Sequential(
            nn.Linear(64*22*22, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 3)
        )

        self.fc2 = nn.Sequential(
            nn.Linear(64*22*22, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 3)
        )

    def forward(self, input1, input2):
        output1 = self.cnn(input1)
        output1 = output1.view(output1.size()[0], -1)
        output1 = self.fc1(output1)

        output2 = self.cnn(input2)
        output2 = output2.view(output2.size()[0], -1)
        output2 = self.fc2(output2)

        return output1, output2
# 创建孪生网络模型实例
siamese_net = SiameseNetwork()
#%%
# 自定义数据集类
from PIL import Image

class CustomDataset(Dataset):
    

    def __getitem__(self, index):
        # 从路径列表中获取图像文件
        img1_path = os.path.join(self.root_dir, self.image_files[index])
        img2_path = os.path.join(self.root_dir, self.image_files[index + 1])

        # 使用PIL库打开图像文件，如果遇到非图像文件则跳过
        try:
            img1 = Image.open(img1_path).convert('L')  # 转换为灰度图像
            img2 = Image.open(img2_path).convert('L')  # 转换为灰度图像
        except Exception as e:
            print(f"Error loading image: {e}. Skipping this image.")
            return self.__getitem__(index + 2)  # 跳过当前图片，直接加载下一张

        # 其他数据处理操作...
        
        # 返回处理后的图像数据
        return img1, img2, label
        

#%%
import os       
'''    
# 设置根目录和数据对
root_dir = 'D:\luanshengNT\KinFaceW-I\images\father-dau'
data_pairs = [...]  # 包含数据路径和对应标签的列表
'''
# 设置根目录
root_dir = r'D:\luanshengNT\KinFaceW-I\images\father-dau'

# 从文件夹中获取所有图像文件名
image_files = sorted(os.listdir(root_dir))

# 创建数据对列表
data_pairs = []
for i in range(0, len(image_files)-1 , 2):
    img1_path = os.path.join(root_dir, image_files[i])
    img2_path = os.path.join(root_dir, image_files[i+1])
    data_pairs.append((img1_path, img2_path, 1))  # 将相似图对应的 label 设置为1

# 打印数据对列表，检查是否正确生成
#  print(data_pairs)



#%%
# 定义数据转换
transform = transforms.Compose([
    transforms.Resize((100, 100)),  # 重置图片大小为100x100
    transforms.ToTensor()  # 转换为张量
])

# 创建自定义数据集实例
custom_dataset = CustomDataset(data_pairs, transform=transform)

# 创建数据加载器
data_loader = DataLoader(custom_dataset, batch_size=32, shuffle=True)

# 实例化孪生网络模型
siamese_net = SiameseNetwork()

# 定义损失函数和优化器
criterion = nn.TripletMarginLoss()
optimizer = optim.Adam(siamese_net.parameters(), lr=0.001)

#定义训练数量
num_epochs = 15
#判断是否为图片类型
from torchvision.utils import save_image
def is_image_file(filename):
    try:
        img = Image.open(filename)  # 尝试打开文件
        img.verify()  # 验证文件是否为图像文件
        img.close()
        return True  # 如果能够成功打开并验证，则为图像文件
    except (IOError, SyntaxError):
        return False  # 如果出现异常，则不是图像文件
# 训练模型
for epoch in range(num_epochs):
    for batch_idx, (input1, input2, label) in enumerate(data_loader):
        optimizer.zero_grad()

        valid_inputs1 = []
        valid_inputs2 = []
        valid_labels = []
        for i in range(len(input1)):
            # 将张量保存为实际图像文件
            save_image(input1[i], 'input1_temp.png')
            save_image(input2[i], 'input2_temp.png')
            
            if is_image_file('input1_temp.png') and is_image_file('input2_temp.png'):
                valid_inputs1.append(input1[i])
                valid_inputs2.append(input2[i])
                valid_labels.append(label[i])

        if not valid_inputs1 or not valid_inputs2:
            continue

        output1, output2 = siamese_net(valid_inputs1, valid_inputs2)
        loss = criterion(output1, output2, valid_labels)
        loss.backward()
        optimizer.step()

#每训练10个batch后打印一次当前的训练状态，其中包括当前的epoch数、batch数以及损失值。
        if batch_idx % 10 == 0:
            print(f'Epoch {epoch+1}, Batch {batch_idx}, Loss: {loss.item()}')